# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: leon

import os
import datetime
import ujson as json
import functools
import logging

from coin.base.book.types import (
    TradeSide,
    create_market_trade_event,
    create_market_ticker_event,
    create_market_instrument_event,
)
from coin.exchange.base.ws.handler import WebsocketHandler
from coin.exchange.okex_futures.book_builder import BookBuilder
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.exchange.base.ws.trade_aggregator import TradeAggregator
from coin.exchange.okex_futures.ws.limit_price_checker.all_products \
    import LimitPriceChecker


class OkexV3HandlerBase(WebsocketHandler):
  BookBuilderType = None
  TraitsType = None

  def __init__(self,
               *,
               event_dist=None,
               feed_checker_config=None,
               agg_view=True,
               check_limit_price=False):
    super().__init__(event_dist=event_dist, feed_checker_config=feed_checker_config)
    self._trade_processor = TradeAggregator(handler=self, agg_view=agg_view)
    self.depth_exch_timestr = {}
    self.trade_exch_timestr = {}
    if check_limit_price:
      self._limit_price_checker = LimitPriceChecker()
    else:
      self._limit_price_checker = None

  @property
  def limit_price_checker(self):
    return self._limit_price_checker

  def add_check_product(self, product):
    if self._limit_price_checker is not None:
      self._limit_price_checker.add_product(product)
    super().add_check_product(product)

  def on_first_book(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict):
      message = json.loads(message.tobytes())

    data = message['data']
    for depth_msg in data:
      instrument_id = depth_msg['instrument_id']
      product = self.TraitsType.get_product_from_instrument_id(instrument_id, record.timestamp)
      symbol = product.native_symbol_v3
      if not self.is_subscribed('book', symbol):
        continue

      current_datetime = datetime.datetime.fromtimestamp(record.timestamp / 1.e9)
      book = self.BookBuilderType(symbol,
                                  current_datetime,
                                  book_builder_observer=self._book_builder_observer)
      bids = depth_msg.get('bids', [])
      asks = depth_msg.get('asks', [])
      book.update_fullbook(bids, asks)
      book.update_timestamp(record.timestamp)
      book.update_topic_id(record.topic_id)
      self._books[symbol] = book

      self.ready = True
      if self._feed_checker:
        self._feed_checker.on_book(product=product, book=book, timestamp=record.timestamp)
      if self._limit_price_checker is not None:
        self._limit_price_checker.on_book(product=product, book=book, timestamp=record.timestamp)
      self.publish_event('book', book.symbol, book)

  def on_diff(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict):
      message = json.loads(message.tobytes())

    data = message['data']
    for depth_msg in data:
      instrument_id = depth_msg['instrument_id']
      product = self.TraitsType.get_product_from_instrument_id(instrument_id, record.timestamp)
      symbol = product.native_symbol_v3
      if not self.is_subscribed('book', symbol):
        continue

      if symbol in self.depth_exch_timestr and self.depth_exch_timestr[symbol] > depth_msg[
          'timestamp']:
        continue
      self.depth_exch_timestr[symbol] = depth_msg['timestamp']

      if symbol not in self._books:
        current_datetime = datetime.datetime.fromtimestamp(record.timestamp / 1.e9)
        book = self.BookBuilderType(symbol,
                                    current_datetime,
                                    book_builder_observer=self._book_builder_observer)
        bids = depth_msg.get('bids', [])
        asks = depth_msg.get('asks', [])
        book.update_fullbook(bids, asks)
        book.update_timestamp(record.timestamp)
        book.update_topic_id(record.topic_id)
        self._books[symbol] = book
      else:
        book = self._books[symbol]
        bids = depth_msg.get('bids', [])
        asks = depth_msg.get('asks', [])
        book.update_diff(bids, asks)
        book.update_timestamp(record.timestamp)
        book.update_topic_id(record.topic_id)

      self.ready = True

      if self._feed_checker:
        self._feed_checker.on_book(product=product, book=book, timestamp=record.timestamp)
      if self._limit_price_checker is not None:
        self._limit_price_checker.on_book(product=product, book=book, timestamp=record.timestamp)
      self.publish_event('book', book.symbol, book)

  def on_coin2(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict):
      message = json.loads(message.tobytes())

    if 'table' not in message:
      return

    if 'depth' in message['table']:
      data = message['data']
      for depth_msg in data:
        instrument_id = depth_msg['instrument_id']
        product = self.TraitsType.get_product_from_instrument_id(instrument_id, record.timestamp)
        symbol = product.native_symbol_v3
        if not self.is_subscribed('book', symbol):
          continue

        if symbol not in self._books:
          current_datetime = datetime.datetime.fromtimestamp(record.timestamp / 1.e9)
          book = self.BookBuilderType(symbol,
                                      current_datetime,
                                      book_builder_observer=self._book_builder_observer)
          self._books[symbol] = book
        else:
          book = self._books[symbol]

        bids = depth_msg.get('bids', [])
        asks = depth_msg.get('asks', [])
        book.update_diff(bids, asks)
        book.update_timestamp(record.timestamp)
        book.update_topic_id(record.topic_id)
        if self._feed_checker:
          self._feed_checker.on_book(product=product, book=book, timestamp=record.timestamp)
        if self._limit_price_checker is not None:
          self._limit_price_checker.on_book(product=product, book=book, timestamp=record.timestamp)
        self.publish_event('book', book.symbol, book)
    elif 'trade' in message['table']:
      self.on_trade(record, queue_data, topic_data)

  def on_trade(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict):
      message = json.loads(message.tobytes())

    self._trade_processor.on_begin_trade()

    data = message['data']
    last_idx = len(data) - 1
    for i, trade_msg in enumerate(data):

      instrument_id = trade_msg['instrument_id']
      product = self.TraitsType.get_product_from_instrument_id(instrument_id, record.timestamp)
      symbol = product.native_symbol_v3

      if not self.is_subscribed('trade', symbol):
        continue

      if symbol in self.trade_exch_timestr and self.trade_exch_timestr[symbol] > trade_msg[
          'timestamp']:
        continue
      self.trade_exch_timestr[symbol] = trade_msg['timestamp']

      price = float(trade_msg['price'])
      qty = self.TraitsType.get_qty_from_trade(trade_msg, record.timestamp)
      side = trade_msg['side']
      if side in ['buy', '2']:
        side = TradeSide.TRADE_BUY_SIDE
      elif side in ['sell', '1']:
        side = TradeSide.TRADE_SELL_SIDE
      else:
        raise ValueError(data)

      trade_event = create_market_trade_event(record.timestamp,
                                              symbol,
                                              price,
                                              qty,
                                              side,
                                              has_more_trade=(i < last_idx),
                                              topic_id=record.topic_id)
      self._trade_processor.on_mid_trade(trade_event)
      if self._feed_checker:
        self._feed_checker.on_trade(product=product,
                                    trade_event=trade_event,
                                    timestamp=record.timestamp)

    self._trade_processor.on_end_trade()

  def on_ticker(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict):
      message = json.loads(message.tobytes())

    data = message['data']
    for ticker_msg in data:
      instrument_id = ticker_msg['instrument_id']
      product = self.TraitsType.get_product_from_instrument_id(instrument_id, record.timestamp)
      symbol = product.native_symbol_v3
      if not self.is_subscribed('ticker', symbol):
        continue

      ticker_event = create_market_ticker_event(timestamp=record.timestamp,
                                                product=product,
                                                low=float(ticker_msg['low_24h']),
                                                high=float(ticker_msg['high_24h']),
                                                vol_24h=float(ticker_msg['volume_24h']))
      self.publish_event('ticker', symbol, ticker_event)

  def on_price_range(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict):
      message = json.loads(message.tobytes())

    data = message['data']
    for price_range_msg in data:
      instrument_id = price_range_msg['instrument_id']
      product = self.TraitsType.get_product_from_instrument_id(instrument_id, record.timestamp)
      ticker_event = create_market_ticker_event(timestamp=record.timestamp,
                                                product=product,
                                                limit_low=float(price_range_msg['lowest']),
                                                limit_high=float(price_range_msg['highest']))
      if self._limit_price_checker is not None:
        self._limit_price_checker.on_ticker(product=product,
                                            ticker_event=ticker_event,
                                            timestamp=record.timestamp)
      symbol = product.native_symbol_v3
      if not self.is_subscribed('ticker', symbol):
        continue
      self.publish_event('ticker', symbol, ticker_event)

  def on_funding_rate(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict):
      message = json.loads(message.tobytes())

    data = message['data']
    for funding_rate_msg in data:
      instrument_id = funding_rate_msg['instrument_id']
      product = self.TraitsType.get_product_from_instrument_id(instrument_id, record.timestamp)
      symbol = product.native_symbol_v3
      if not self.is_subscribed('instrument', symbol):
        continue

      instrument_event = create_market_instrument_event(timestamp=record.timestamp,
                                                        product=product,
                                                        funding_rate=float(
                                                            funding_rate_msg['funding_rate']),
                                                        topic_id=record.topic_id)
      self.publish_event('instrument', symbol, instrument_event)

  def on_mark_price(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict):
      message = json.loads(message.tobytes())

    data = message['data']
    for mark_price_msg in data:
      if not isinstance(mark_price_msg, dict):
        continue
      instrument_id = mark_price_msg['instrument_id']
      product = self.TraitsType.get_product_from_instrument_id(instrument_id, record.timestamp)
      symbol = product.native_symbol_v3
      if not self.is_subscribed('instrument', symbol):
        continue

      instrument_event = create_market_instrument_event(timestamp=record.timestamp,
                                                        product=product,
                                                        mark_price=float(
                                                            mark_price_msg['mark_price']),
                                                        topic_id=record.topic_id)
      self.publish_event('instrument', symbol, instrument_event)

  def on_open_interest(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict):
      message = json.loads(message.tobytes())

    if 'instrument_id' not in message:
      logging.error(message)
      return

    instrument_id = message['instrument_id']
    product = self.TraitsType.get_product_from_instrument_id(instrument_id, record.timestamp)
    symbol = product.native_symbol_v3
    if not self.is_subscribed('instrument', symbol):
      return

    instrument_event = create_market_instrument_event(timestamp=record.timestamp,
                                                      product=product,
                                                      open_interest=float(message['amount']),
                                                      topic_id=record.topic_id)
    self.publish_event('instrument', symbol, instrument_event)


OkexSpotV3HandlerBase = OkexV3HandlerBase

if ('RUST_FEED_HANDLER' in os.environ) and os.environ['RUST_FEED_HANDLER'] == '1':
  try:
    from native_rust import feed_handler
    print("okex base: using rust feed handler")

    class OkexV3RustHandlerBase(OkexV3HandlerBase):
      BookBuilderType = None
      TraitsType = None

      def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # TODO(jhkim): spot == futures?
        self._rust_handler = feed_handler.OkexV3FuturesHandler()
        # self.publish_event will use _event_dist.publish
        # to register to _event_dist, use self.subscribe_trade
        self._rust_handler.register_book_callback(functools.partial(self.publish_event, 'book'))
        self._rust_handler.register_trade_callback(functools.partial(self.publish_event, 'trade'))

      def post_subscribe(self, event_type, symbol, callback):
        assert type(symbol) == str
        if event_type == 'book':
          self._rust_handler.register_book_symbol(symbol)
        elif event_type == 'trade':
          self._rust_handler.register_trade_symbol(symbol)
        elif event_type == 'instrument':
          pass
        else:
          raise RuntimeError('Okex v3 rust handler not implemented for event type: ', event_type)

      def on_first_book(self, record, queue_data, topic_data):
        if isinstance(record.data, (dict, list)):
          record_bytes = json.dumps(record.data).encode()
        else:
          record_bytes = record.data.tobytes()

        self._rust_handler.on_snapshot(record.timestamp, record_bytes, {}, {})

      def on_diff(self, record, queue_data, topic_data):
        if isinstance(record.data, (dict, list)):
          record_bytes = json.dumps(record.data).encode()
        else:
          record_bytes = record.data.tobytes()
        self._rust_handler.on_diff(record.timestamp, record_bytes, {}, {})

      def on_coin2(self, record, queue_data, topic_data):
        if isinstance(record.data, (dict, list)):
          record_bytes = json.dumps(record.data).encode()
        else:
          record_bytes = record.data.tobytes()
        self._rust_handler.on_coin2(record.timestamp, record_bytes, {}, {})

    OkexV3HandlerBase = OkexV3RustHandlerBase

  except ImportError:
    import logging
    logging.warning('Unable to import rust feed_handler.')
